import mindspore as ms
import bert_model as bert
import mindspore.nn as nn
import mindspore.ops as ops




class classifier(nn.Cell):
    def __init__(self,config, is_training,fc_dim1,class_num,bert_weight_path) :
        super().__init__()
        self.bert=bert.BertModel(config, is_training)
        ms.load_checkpoint(bert_weight_path,self.bert)
        print('init model successfully')
        self.fc1=nn.Dense(self.bert.embedding_size,class_num)
        self.fc2=nn.Dense(fc_dim1,class_num)
        #self.relu1=ops.ReLU()
        self.dropout = nn.Dropout(1 - 0.1)

    def construct(self,text_ids, segment_ids, mask_ids):
        _, pooler_output, _=self.bert(text_ids, segment_ids, mask_ids)
        fc1_out=self.dropout(self.fc1(pooler_output))
        #fc2_out=self.fc2(fc1_out)
        return fc1_out

class classifier_predict(nn.Cell):
    def __init__(self,config, is_training,fc_dim1,class_num,bert_weight_path) :
        super().__init__()
        self.bert=bert.BertModel(config, is_training)
        self.fc1=nn.Dense(self.bert.embedding_size,class_num)
        self.fc2=nn.Dense(fc_dim1,class_num)
        #self.relu1=ops.ReLU()
        self.dropout = nn.Dropout(1 - 0.1)

    def construct(self,text_ids, segment_ids, mask_ids):
        _, pooler_output, _=self.bert(text_ids, segment_ids, mask_ids)
        fc1_out=self.dropout(self.fc1(pooler_output))
        #fc2_out=self.fc2(fc1_out)
        return fc1_out

#进一步封装model，包含loss function ，使得输出直接为loss
class MyWithLossCell(nn.Cell):
   def __init__(self, backbone, loss_fn):
       super(MyWithLossCell, self).__init__(auto_prefix=False)
       self._backbone = backbone
       self._loss_fn = loss_fn

   def construct(self, x, y, z, label):
       out = self._backbone(x, y, z)
       return self._loss_fn(out, label)

   @property
   def backbone_network(self):
       return self._backbone